import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR10
from torch.utils.data import DataLoader
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score, adjusted_rand_score, silhouette_samples
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import numpy as np

# ================= 可视化函数 =================

def visualize_misclustered_images(features, labels, preds, images, class_names=None, max_show=25):
    mismatches = (labels != preds)
    mismatched_idx = np.where(mismatches)[0]
    print(f"总共发现 {len(mismatched_idx)} 个预测错误样本")

    plt.figure(figsize=(10, 10))
    for i, idx in enumerate(mismatched_idx[:max_show]):
        plt.subplot(5, 5, i+1)
        img = images[idx].transpose(1, 2, 0)  # C×H×W to H×W×C
        plt.imshow(img)
        plt.axis('off')
        title = f"T:{labels[idx]} P:{preds[idx]}"
        plt.title(title, fontsize=8)
    plt.suptitle("Misclustered Samples (True vs Pred)", fontsize=12)
    plt.tight_layout()
    plt.show()

def visualize_dual_tsne(features, labels, preds):
    tsne = TSNE(n_components=2, random_state=42)
    reduced = tsne.fit_transform(features)

    plt.figure(figsize=(12, 5))
    plt.subplot(1, 2, 1)
    plt.scatter(reduced[:, 0], reduced[:, 1], c=labels, cmap='tab10', s=5)
    plt.title("t-SNE by True Labels")
    plt.colorbar()

    plt.subplot(1, 2, 2)
    plt.scatter(reduced[:, 0], reduced[:, 1], c=preds, cmap='tab10', s=5)
    plt.title("t-SNE by Cluster Labels")
    plt.colorbar()

    plt.tight_layout()
    plt.show()

def visualize_silhouette_distribution(features, preds):
    sil_samples = silhouette_samples(features, preds)
    plt.figure(figsize=(8, 4))
    plt.hist(sil_samples, bins=50, color='skyblue', edgecolor='k')
    plt.title("Silhouette Coefficient Distribution")
    plt.xlabel("Silhouette Coefficient")
    plt.ylabel("Frequency")
    plt.grid(True)
    plt.show()

def visualize_tsne(features, labels, title="t-SNE Visualization"):
    tsne = TSNE(n_components=2, random_state=42)
    reduced = tsne.fit_transform(features)
    plt.figure(figsize=(8,6))
    plt.scatter(reduced[:,0], reduced[:,1], c=labels, cmap='tab10', s=5)
    plt.title(title)
    plt.colorbar()
    plt.show()

# ================= AutoEncoder 适配 CIFAR10 =================
class AutoEncoder(nn.Module):
    def __init__(self, latent_dim=64):
        super(AutoEncoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),              # 3x32x32 = 3072
            nn.Linear(3072, 1024),
            nn.ReLU(),
            nn.Linear(1024, latent_dim)
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 1024),
            nn.ReLU(),
            nn.Linear(1024, 3072),
            nn.Sigmoid(),
            nn.Unflatten(1, (3, 32, 32))
        )

    def forward(self, x):
        z = self.encoder(x)
        out = self.decoder(z)
        return out

# ================= 数据加载函数 =================
def load_cifar10():
    transform = transforms.Compose([transforms.ToTensor()])
    train = CIFAR10(root="./data", train=True, download=True, transform=transform)
    test = CIFAR10(root="./data", train=False, download=True, transform=transform)
    return train, test

# ================= AE训练函数 =================
def train_autoencoder(model, dataset, epochs=10):
    dataloader = DataLoader(dataset, batch_size=128, shuffle=True)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.MSELoss()
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for x, _ in dataloader:
            x = x.cuda()
            out = model(x)
            loss = criterion(out, x)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f"Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}")

# ================= 特征提取函数 =================
def extract_features(model, dataset):
    dataloader = DataLoader(dataset, batch_size=128)
    model.eval()
    features, labels = [], []
    with torch.no_grad():
        for x, y in dataloader:
            x = x.cuda()
            z = model.encoder(x).cpu()
            features.append(z)
            labels.append(y)
    return torch.cat(features).numpy(), torch.cat(labels).numpy()

# ================= 聚类评估函数 =================
def cluster_and_evaluate(features, labels):
    kmeans = KMeans(n_clusters=10, random_state=42)
    preds = kmeans.fit_predict(features)
    sil = silhouette_score(features, preds)
    ari = adjusted_rand_score(labels, preds)
    print(f"Silhouette Score: {sil:.4f}, ARI: {ari:.4f}")
    return preds

# ================= 主流程 =================
def main():
    train_set, test_set = load_cifar10()

    for latent_dim in [16, 32, 64, 128]:
        print(f"\nTesting latent_dim = {latent_dim}")
        model = AutoEncoder(latent_dim=latent_dim).cuda()
        train_autoencoder(model, train_set, epochs=5)
        features, labels = extract_features(model, test_set)

        # 聚类
        preds = cluster_and_evaluate(features, labels)

        # 可视化
        visualize_tsne(features, preds, title=f"AE latent_dim={latent_dim} (Cluster Labels)")
        visualize_dual_tsne(features, labels, preds)
        visualize_silhouette_distribution(features, preds)

        # 图像收集用于错误聚类展示
        test_loader = DataLoader(test_set, batch_size=128)
        all_images = []
        for x, _ in test_loader:
            all_images.append(x)
        all_images = torch.cat(all_images).numpy()
        visualize_misclustered_images(features, labels, preds, all_images)

    # 原始图像 baseline
    print("\nBaseline: 使用原始图像")
    dataloader = DataLoader(test_set, batch_size=128)
    all_x, all_y = [], []
    for x, y in dataloader:
        all_x.append(x.view(x.size(0), -1))
        all_y.append(y)
    x_raw = torch.cat(all_x).numpy()
    y_raw = torch.cat(all_y).numpy()
    preds = cluster_and_evaluate(x_raw, y_raw)
    visualize_tsne(x_raw, preds, title="Raw Input t-SNE")

if __name__ == '__main__':
    main()
